# prediction_model.py - 预测模型构建
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (classification_report, confusion_matrix,
                             roc_auc_score, precision_recall_curve, roc_curve)
from sklearn.preprocessing import LabelEncoder
from config import MODEL_CONFIG


class CouponPredictionModel:
    def __init__(self, df, user_features):
        self.df = df
        self.user_features = user_features
        self.model = None
        self.feature_names = None
        self.X_test = None
        self.y_test = None

    def prepare_training_data(self):
        """准备训练数据"""
        print("\n=== 准备训练数据 ===")

        # 创建目标变量 - 基于用户历史优惠券使用行为
        user_coupon_behavior = self.df[self.df['使用状态'] != -1].groupby('用户ID').agg({
            '使用状态': 'mean'  # 用户历史优惠券使用率
        }).reset_index()

        # 定义目标变量：使用率超过50%的用户标记为正样本
        user_coupon_behavior['should_send_coupon'] = (user_coupon_behavior['使用状态'] > 0.3).astype(int)

        print(f"目标变量分布:")
        print(user_coupon_behavior['should_send_coupon'].value_counts())

        # 准备特征数据
        X = self.user_features.select_dtypes(include=[np.number]).copy()

        # 移除可能泄露目标变量的特征
        columns_to_exclude = [
            '客户分群', '优惠券使用率', '总消费额', '购买频次',
            '最近购买间隔', '客户生命周期'
        ]
        X = X.drop(columns=[col for col in columns_to_exclude if col in X.columns], errors='ignore')

        # 设置目标变量
        y = user_coupon_behavior.set_index('用户ID')['should_send_coupon']

        # 对齐数据
        common_users = X.index.intersection(y.index)
        X = X.loc[common_users]
        y = y.loc[common_users]

        # 处理缺失值
        X = X.fillna(X.median())

        print(f"训练数据维度: X {X.shape}, y {y.shape}")
        print(f"正样本比例: {y.mean():.2%}")

        return X, y

    def train_model(self, X, y):
        """训练预测模型"""
        print("\n=== 训练预测模型 ===")

        # 划分训练测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y,
            test_size=MODEL_CONFIG['test_size'],
            random_state=MODEL_CONFIG['random_state'],
            stratify=y
        )

        # 训练随机森林模型
        self.model = RandomForestClassifier(
            n_estimators=MODEL_CONFIG['n_estimators'],
            random_state=MODEL_CONFIG['random_state'],
            class_weight='balanced'
        )
        self.model.fit(X_train, y_train)

        self.feature_names = X.columns.tolist()
        self.X_test = X_test
        self.y_test = y_test

        print("模型训练完成")
        return X_train, X_test, y_train, y_test

    def evaluate_model(self):
        """评估模型性能"""
        if self.model is None or self.X_test is None:
            print("请先训练模型")
            return

        print("\n=== 模型性能评估 ===")

        # 预测
        y_pred = self.model.predict(self.X_test)
        y_pred_proba = self.model.predict_proba(self.X_test)[:, 1]

        # 分类报告
        print("\n分类报告:")
        print(classification_report(self.y_test, y_pred))

        # AUC Score
        auc_score = roc_auc_score(self.y_test, y_pred_proba)
        print(f"AUC Score: {auc_score:.4f}")

        # 可视化评估结果
        self._plot_model_evaluation(y_pred, y_pred_proba)

        return {
            'y_pred': y_pred,
            'y_pred_proba': y_pred_proba,
            'auc_score': auc_score
        }

    def _plot_model_evaluation(self, y_pred, y_pred_proba):
        """绘制模型评估图表"""
        fig = plt.figure(figsize=(16, 12))

        # 1. 特征重要性
        ax1 = plt.subplot(2, 2, 1)
        feature_importance = pd.DataFrame({
            'feature': self.feature_names,
            'importance': self.model.feature_importances_
        }).sort_values('importance', ascending=False)

        plt.barh(feature_importance['feature'][:15], feature_importance['importance'][:15])
        plt.title('特征重要性TOP15')
        plt.xlabel('重要性')

        # 2. 混淆矩阵
        ax2 = plt.subplot(2, 2, 2)
        cm = confusion_matrix(self.y_test, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=['不发放', '发放'],
                    yticklabels=['不发放', '发放'])
        plt.title('混淆矩阵')
        plt.ylabel('真实标签')
        plt.xlabel('预测标签')

        # 3. ROC曲线
        ax3 = plt.subplot(2, 2, 3)
        fpr, tpr, _ = roc_curve(self.y_test, y_pred_proba)
        plt.plot(fpr, tpr, color='darkorange', lw=2,
                 label=f'ROC curve (AUC = {roc_auc_score(self.y_test, y_pred_proba):.2f})')
        plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('ROC曲线')
        plt.legend(loc="lower right")

        # 4. 精确率-召回率曲线
        ax4 = plt.subplot(2, 2, 4)
        precision, recall, _ = precision_recall_curve(self.y_test, y_pred_proba)
        plt.plot(recall, precision, color='blue', lw=2)
        plt.xlabel('召回率')
        plt.ylabel('精确率')
        plt.title('P-R曲线')

        plt.tight_layout()
        plt.show()

    def get_feature_importance(self, top_n=15):
        """获取特征重要性"""
        if self.model is None:
            print("请先训练模型")
            return None

        importance_df = pd.DataFrame({
            'feature': self.feature_names,
            'importance': self.model.feature_importances_
        }).sort_values('importance', ascending=False)

        return importance_df.head(top_n)

    def predict_coupon_propensity(self, user_data):
        """预测用户优惠券使用倾向"""
        if self.model is None:
            print("请先训练模型")
            return None

        # 准备特征数据
        X = user_data[self.feature_names].fillna(user_data[self.feature_names].median())

        # 预测概率
        propensity_scores = self.model.predict_proba(X)[:, 1]

        # 根据概率分箱
        propensity_levels = pd.cut(propensity_scores,
                                   bins=[0, 0.3, 0.6, 0.8, 1.0],
                                   labels=['低倾向', '中倾向', '高倾向', '极高倾向'])

        result_df = user_data.copy()
        result_df['优惠券使用倾向'] = propensity_scores
        result_df['倾向等级'] = propensity_levels

        return result_df

    def get_model_info(self):
        """获取模型信息"""
        if self.model is None:
            return "模型未训练"

        info = {
            'model_type': type(self.model).__name__,
            'n_features': len(self.feature_names),
            'feature_names': self.feature_names
        }

        return info